package io.github.javpower.knowledgerag.service;

import io.github.javpower.knowledgerag.config.RagConfig;
import io.github.javpower.knowledgerag.core.ContextBuilder;
import io.github.javpower.knowledgerag.core.DocumentReader;
import io.github.javpower.knowledgerag.core.Vectorizer;
import io.github.javpower.knowledgerag.req.QuestionReq;
import io.github.javpower.vectorexclient.VectorRexClient;
import io.github.javpower.vectorexclient.builder.QueryBuilder;
import io.github.javpower.vectorexclient.entity.MetricType;
import io.github.javpower.vectorexclient.entity.ScalarField;
import io.github.javpower.vectorexclient.entity.VectorFiled;
import io.github.javpower.vectorexclient.req.CollectionDataAddReq;
import io.github.javpower.vectorexclient.req.VectoRexCollectionReq;
import io.github.javpower.vectorexclient.res.ServerResponse;
import io.github.javpower.vectorexclient.res.VectorSearchResult;
import io.github.javpower.vectorexclient.util.GsonUtil;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

@Service
public class QnAService {

    private final String collectionName;
    private final VectorRexClient client;
    private final String OLLAMA_API_URL;
    private final String document;
    private final String model;

    @Autowired
    public QnAService(RagConfig config) {
        collectionName=config.getCollectionName();
        OLLAMA_API_URL=config.getOllamaApiUrl();
        document=config.getDocument();
        model=config.getModel();
        this.client =  new VectorRexClient(config.getVectorUrl(), config.getVectorUser(), config.getVectorPassword());
        initialize();
    }
    private void initialize() {
        boolean anyMatch = client.getCollections().getData().stream().anyMatch(collection -> collection.getCollectionName().equals(collectionName));
        if (anyMatch) {
            return;
        }
        List<String> chunks = new ArrayList<>();
        List<List<Float>> vectors = new ArrayList<>();

        // 读取文档并分块
        chunks.addAll(DocumentReader.readAndChunkDocument(document, 1024));

        // 向量化
        vectors.addAll(Vectorizer.vectorizeChunks(OLLAMA_API_URL,model,chunks));

        // 创建集合
        createCollection(collectionName, vectors.get(0).size());

        // 添加数据
        for (int i = 0; i < chunks.size(); i++) {
            addData(collectionName,i,chunks.get(i), vectors.get(i));
        }
    }

    private void createCollection(String collectionName, int vectorDimensions) {

        List<ScalarField> scalarFields = new ArrayList<>();
        scalarFields.add(ScalarField.builder().name("id").isPrimaryKey(true).build());
        scalarFields.add(ScalarField.builder().name("text").isPrimaryKey(false).build());

        List<VectorFiled> vectorFields = new ArrayList<>();
        vectorFields.add(VectorFiled.builder().name("vector").metricType(MetricType.FLOAT_COSINE_DISTANCE).dimensions(vectorDimensions).build());

        VectoRexCollectionReq req = VectoRexCollectionReq.builder()
                .collectionName(collectionName)
                .scalarFields(scalarFields)
                .vectorFileds(vectorFields)
                .build();

        ServerResponse<Void> response = client.createCollection(req);
        if (response.isSuccess()) {
            System.out.println("集合创建成功");
        } else {
            System.out.println("集合创建失败：" + response.getMsg());
        }
    }

    private void addData(String collectionName, Integer id, String text, List<Float> vector) {
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("id", id);
        metadata.put("text", text);
        metadata.put("vector", vector);

        CollectionDataAddReq req = CollectionDataAddReq.builder()
                .collectionName(collectionName)
                .metadata(metadata)
                .build();

        ServerResponse<Void> response = client.addCollectionData(req);
        if (response.isSuccess()) {
            System.out.println("数据添加成功");
        } else {
            System.out.println("数据添加失败：" + response.getMsg());
        }
    }

    private List<VectorSearchResult> queryVectorDatabase(String collectionName, List<Float> queryVector, int topK) {
        QueryBuilder builder = QueryBuilder.lambda(collectionName);
        builder.vector("vector", queryVector).topK(topK);
        ServerResponse<List<VectorSearchResult>> response = client.queryCollectionData(builder);
        if (response.isSuccess()) {
            return response.getData();
        } else {
            System.out.println("数据查询失败：" + response.getMsg());
            return Collections.emptyList();
        }
    }

    public String generateAnswer(QuestionReq req) {
        try {
            // 将问题转换为向量
            List<Float> queryVector = Vectorizer.getVector(OLLAMA_API_URL,model,req.getQuestion());

            // 查询向量数据库
            List<VectorSearchResult> results = queryVectorDatabase(collectionName, queryVector, 2);

            // 构建上下文
            String context = ContextBuilder.buildContext(results);

            // 使用模型生成答案
            String answer = executeQwenModel(model,req.getQuestion(),context);

            return answer;
        } catch (Exception e) {
            e.printStackTrace();
            return "Error generating answer";
        }
    }

    private String executeQwenModel(String model,String question, String context) {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        String json = "{\"model\": \""+model+"\", \"prompt\": \"" + question + " " + context + "\"}";
        RequestBody body = RequestBody.create(json, mediaType);
        Request request = new Request.Builder()
                .url(OLLAMA_API_URL+"/generate")
                .post(body)
                .addHeader("Content-Type", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }
            String res = response.body().string();
            String[] split = res.split("\n");
            String answerRes = Arrays.stream(split).map(v-> {
                Map s = GsonUtil.fromJson(v, Map.class);
                return s.get("response").toString();
            }).collect(Collectors.joining(""));
            return answerRes;
        } catch (IOException e) {
            e.printStackTrace();
            return "Error generating answer";
        }
    }
}